open Glob.Bin

module Vector = struct
  let mul v1 v2 = 
      let n = Array.length v1 in
      assert (n = Array.length v2);
      let sum = ref 0. in
      for k = 0 to n - 1 do sum += v1.(k) *. v2.(k) done;
      !sum
(*      Array.fold_left (fun sum v1_i ->
      sum +. v2.(i)*.v1_i
  ) 0. v1*)
  let modul v = sqrt (mul v v)
  module Bin = struct
    let ( |*= ) v const = 
        for i = 0 to Array.length v - 1 do
            v.(i) <- v.(i) *. const
        done
    ;;

    (** copy v1 to v2 *)
    let ( ->| ) v1 v2 =
        Array.iteri (Array.set v2) v1

(*    let ( |* ) v const = 
        Array.map ( ( *. ) const ) v*)
    let ( *| ) const v = 
        Array.map ( ( *. ) const ) v

    let ( -| ) v1 v2 =
        Array.mapi (fun i v1_i -> v1_i -. v2.(i)) v1
    let ( +| ) v1 v2 =
        Array.mapi (fun i v1_i -> v1_i +. v2.(i)) v1
  end
  let select components v = Array.map (fun i -> v.(i)) components
end

module Vect2 = struct
    type vector={x:float;y:float}
    let v x y={x=x;y=y}
    let mul v1 v2 = v1.x *. v2.x +. v1.y *. v2.y
    let md v1 = sqrt (v1.x**2. +. v1.y**2.)
    let angle v1 v2 = acos ((mul v1 v2) /. ((md v1)*.(md v2)))
    let ( * ) v1 v2 = mul v1 v2
end

module Vect3_s : sig
    type t = private V of float array
(*    type t = V of float array*)
    val v : float array -> t
    val components : t -> float array
    val mul : t -> t -> float
    val modulus : t -> float
    val beta_gamma : t -> float*float
    val rev : t -> t
end
    = struct
    type t = V of float array
    let v array = assert (Array.length array = 3);
        V array
    let components (V v) = v
    let mul (V v1) (V v2) = List.fold_left (fun sum i ->
        sum +. v2.(i)*.v1.(i)
    ) 0. [0;1;2]
    let modulus v = sqrt (mul v v)
    let beta_gamma v =
        let beta = modulus v in
        let gamma = 1. /. sqrt ( 1. -. beta*.beta ) in
        (beta,gamma)
    let rev (V v) = V (Array.map (fun c -> -.c) v)
end

module Vect3_arr = struct
    type t = float array
    let x (v:t) = v.(0)
    let y (v:t) = v.(1)
    let z (v:t) = v.(2)
    let perp2 v = v.(0)**2. +. v.(1)**2.
    let perp v = sqrt (perp2 v)
    let mul v1 v2 = List.fold_left (fun sum i ->
        sum +. v2.(i)*.v1.(i)
    ) 0. [0;1;2]
    let modulus v = sqrt (mul v v)
    let beta_gamma v =
        let beta = modulus v in
        let gamma = 1. /. sqrt ( 1. -. beta*.beta ) in
        (beta,gamma)

    let cos_angle v1 v2 = (mul v1 v2) /. ((modulus v1)*.(modulus v2))
    let angle v1 v2 = acos (cos_angle v1 v2)
    let theta v = angle v [| 0.; 0.; 1. |]
    let cos_theta v = cos (theta v)
(*    let tanx v = (x v) /. (modulus v)
    let tany v = (y v) /. (modulus v)
    let tanxy v = let m = modulus v in (x v)/.m, (y v)/.m
    let tan v = (perp v) /. (modulus v)*)
    let tanx_theta v = (x v) /. (z v)
    let tany_theta v = (y v) /. (z v)
    let tanxy_theta v = (x v)/.(z v), (y v)/.(z v)
    let tan_theta v = (perp v) /. (z v)
end

module Make_LVect4 (Order: sig val et_is_on0: bool end) = 
struct
    module Vect3 = Vect3_arr
    open Vector.Bin
    type vector=float array
    let v array = assert (Array.length array = 4);
        array
    let metric_E = [| 1.;  1.;  1.;  1. |]
    let numer4 = [| 0; 1; 2; 3; |]
    let metric_M = 
        if Order.et_is_on0
        then [| 1.; -1.; -1.; -1. |]
        else [| -1.; -1.; -1.; 1. |]
    let v4et, v3beg, v3med, v3end =
        if Order.et_is_on0
        then 0,1,2,3
        else 3,0,1,2
    let v3_components = [| v3beg; v3med; v3end; |]
    let mul metric numer v1 v2 = Array.fold_left (fun sum i ->
        sum +. v2.(i)*.v1.(i)*.metric.(i) 
    ) 0. numer
    let mul_M = mul metric_M numer4
    let mass2 v = mul_M v v 
    let mass v = sqrt ( mul_M v v )
    let vect3 v =
        Vector.select v3_components v 

    let momentum_mod v = sqrt (mul metric_E v3_components v v)
(*    let momentum_mod v = Vect3.modulus (vect3 v)*)
(*    let momentum v = Vector.select v3_components v *)

    let make_ep e p :vector =
        let a = Array.create 4 0. in
        Array.iteri (fun i i4 -> a.(i4) <- p.(i)) v3_components;
        a.(v4et) <- e;
        a

    let make_mp m p :vector =
        let a = Array.create 4 0. in
        Array.iteri (fun i i4 -> a.(i4) <- p.(i)) v3_components;
        a.(v4et) <- sqrt (m**2. +. p.(0)**2. +. p.(1)**2. +. p.(2)**2.);
        a

    let boost_of (v: vector) = 
        let a = Vector.select v3_components v in
        Array.map (fun c -> c /. v.(v4et)) a
    let rev_boost_of (v: vector) = (-1.) *| (boost_of v)
    let beta_gamma (v: vector) = Vect3.beta_gamma (boost_of v)

    let make_boost (v: vector) bv = 
        let beta, gamma = Vect3.beta_gamma bv in
        let result = Array.make 4 0. in
        let bx = Vect3.mul bv (vect3 v) in
        let coef=(gamma-.1.) /. (beta*.beta) in
        Array.iteri (fun i i4 ->
            result.(i4) <- v.(i4) +. coef *. bx *. bv.(i) 
                +. gamma *. v.(v4et) *. bv.(i)
        ) v3_components;
        result.(v4et) <- gamma *. ( v.(v4et) +. bx ) ;
        result
    ;;

    let cos_angle v1 v2 = Vect3.cos_angle (vect3 v1) (vect3 v2)
    let angle v1 v2 = Vect3.angle (vect3 v1) (vect3 v2)
    let theta v = Vect3.theta (vect3 v)
    let cos_theta v = Vect3.cos_theta (vect3 v)
    let tanx_theta = Vect3.tanx_theta $ vect3
    let tany_theta = Vect3.tany_theta $ vect3
    let tanxy_theta = Vect3.tanxy_theta $ vect3
    let tan_theta = Vect3.tan_theta $ vect3

    let en v = v.(v4et)
    let x = Vect3.x $ vect3
    let y = Vect3.y $ vect3
    let z = Vect3.z $ vect3
    let perp = Vect3.perp $ vect3
    let perp2 = Vect3.perp2 $ vect3


   module Bin = struct
     let ( !! ) = momentum_mod
   end
end

(** standard 4-momenta numeration: E,px,py,pz *)
module Lorenz_vector = Make_LVect4(struct let et_is_on0 = true end)
(** 4-momenta numeration accepted CERN ROOT: px,py,pz,E *)
module Lorenz_vector_ROOT = Make_LVect4(struct let et_is_on0 = false end)

